#include "Dwarf.Unmanaged.Mpi.h"

//Init arrays.
Solver :: Solver()
{
    lgrid = 0;
    rgrid = 0;
    result = 0;

    msize = 0;
    size = 0;

    int argc = 0;
    char **argv;
    MPI_Init(&argc, &argv);                             // Init MPI

    MPI_Comm_rank(MPI_COMM_WORLD, &commRank);           // Retrieve rank of the process
    MPI_Comm_size(MPI_COMM_WORLD, &commSize);           // Retrieve count of processes

    isRootThread = commRank == NUMBER_ROOT_PROCESS;
}

//Dispose arrays.
Solver :: ~Solver()
{
    MPI_Finalize();                                     //Dispose MPI

    for (int i = 0; i < size; i ++) 
    {            
        delete lgrid[i];
        delete result[i];
    }

    for (int i = 0; i < msize; i ++) 
    {            
        delete rgrid[i];
    }

    delete lgrid;
    delete rgrid;
    delete result;
}

static void Transpose(double** matrix, double** transpose, int n)
{
    for (int i = 0; i < n; i++)
    {
        transpose[i][i] = matrix[i][i];
        for (int j = i + 1; j < n; j++)
        {
            transpose[i][j] = matrix[j][i];
            transpose[j][i] = matrix[i][j];
        }
    }
}

// MPI based method for matrix-matrix multiplication.
void Solver::solve() 
{    
    int rowNumSend = 0;     //Number of sent rows.
    int rowNumRecv = 0;     //Number of recieved rows.
    MPI_Status status;      //Status of receiving.
    
    if (isRootThread)
    {   
		// Root Process 

        //Give each worker process a row to start with
        for (int i = 1; i < commSize; i ++) 
        {
            MPI_Send(&rowNumSend, 1, MPI_INT, i, i, MPI_COMM_WORLD);

            if (rowNumSend < msize) 
            {
                MPI_Send(lgrid[rowNumSend], msize, MPI_DOUBLE, i, i, MPI_COMM_WORLD);
                rowNumSend ++;
            }
        }

        //Loop until all rows sent and all results received
        while(rowNumRecv < msize) 
        {
            //Receive row results from any process.
            MPI_Probe(MPI_ANY_SOURCE, MPI_ANY_TAG, MPI_COMM_WORLD, &status);

            MPI_Recv(result[status.MPI_TAG], msize, MPI_DOUBLE, status.MPI_SOURCE, MPI_ANY_TAG, MPI_COMM_WORLD, &status);
            rowNumRecv ++;

            //Send next row (if any) to process  that completed.
            //Row number out of range is all done signal
            MPI_Send(&rowNumSend, 1, MPI_INT, status.MPI_SOURCE, status.MPI_SOURCE, MPI_COMM_WORLD);
            if (rowNumSend < msize) 
            {
                MPI_Send(lgrid[rowNumSend], msize, MPI_DOUBLE, status.MPI_SOURCE, status.MPI_SOURCE, MPI_COMM_WORLD);
                rowNumSend ++;
            }
        }        
    }
    else
    {   
		// Worker Processes

		// transpose the right hand array to improve memory accessing
		double** rtransposed = new double*[msize];
		for (int i = 0; i < msize; i++)
		{
			rtransposed[i] = new double[msize];
		}
		Transpose(rgrid, rtransposed, msize);

        // temp work rows 
        double* row = new double[msize];
		double* leftRow = new double[msize];

        //Process rows until told to stop.
        while(true) 
        {
            // get row number to process
			MPI_Recv(&rowNumSend, 1, MPI_INT, NUMBER_ROOT_PROCESS, MPI_ANY_TAG, MPI_COMM_WORLD, &status);
            if (rowNumSend >= msize) break;

            //Receive row.
            MPI_Recv(leftRow, msize, MPI_DOUBLE, NUMBER_ROOT_PROCESS, MPI_ANY_TAG, MPI_COMM_WORLD, &status);

			//Multiple leftRow by rgrid (transposed)
			//    Note that row and cell are factored out
			//    since the optimizers shouldn't hoist all of 
			//    the common addressing subexpressions without doing it. 
			//    (Strickly speaking it would be incorrect to do.)
			for (int i = 0; i < msize; i ++ )
            {
				double cell = 0.0;
				double* rightColumn = rtransposed[i];
                for (int j = 0; j < msize; j ++ )
                {
					cell += leftRow[j] * rightColumn[j];
				}
				row[i] = cell;
            }

            //Return resulted row to root process.
            MPI_Send(row, msize, MPI_DOUBLE, NUMBER_ROOT_PROCESS, rowNumSend, MPI_COMM_WORLD); 
		}

		//delete array temps
    	for(int i=0; i<msize; i++)
		{
			delete rtransposed[i];
		}
		delete rtransposed;
		delete row;
		delete leftRow;
	}
}